import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import roc_curve, auc, confusion_matrix, accuracy_score, f1_score, precision_score, recall_score
import seaborn as sns
from utils.utils import *

def bootstrap_auc_ci(model, X, y, n_bootstraps=200, ci=0.95, random_state=None):
    rng = np.random.RandomState(random_state)
    aucs = []

    for _ in range(n_bootstraps):
        indices = rng.randint(0, len(y), len(y))
        X_sample = X[indices] if isinstance(X, np.ndarray) else X.iloc[indices]
        y_sample = y[indices] if isinstance(y, np.ndarray) else y.iloc[indices]

        prob_sample = model.predict_proba(X_sample)[:, 1]
        fpr_sample, tpr_sample, _ = roc_curve(y_sample, prob_sample)
        aucs.append(auc(fpr_sample, tpr_sample))

    lower = np.percentile(aucs, (1 - ci) / 2 * 100)
    upper = np.percentile(aucs, (1 + ci) / 2 * 100)
    return lower, upper


def plot_auc_cm(result: dict,
                cm=None,
                draw_ci=False,
                auc_title="ROC Curve",
                cm_title="Confusion Matrix",
                save_path=None,
                filename="auc_cm.pdf"):
    if cm is not None:
        fig, axes = plt.subplots(1, 2, figsize=(14, 6))
        ax1, ax2 = axes[0], axes[1]
    else:
        fig, ax1 = plt.subplots(figsize=(8, 6))
        ax2 = None

    # === 左侧：ROC 曲线 ===
    for name, info in result.items():
        mean_fpr = info["fpr"]
        mean_tpr = info["tpr"]
        auc_score = info["auc"]
        ci = info.get("ci", None)
        label = f"{name} (AUC = {auc_score:.3f}"
        if ci:
            label += f", 95% CI [{ci[0]:.3f}, {ci[1]:.3f}]"
        label += ")"
        ax1.plot(mean_fpr, mean_tpr, label=label)

        if draw_ci and "tpr_ci_lower" in info and "tpr_ci_upper" in info:
            ax1.fill_between(mean_fpr, info["tpr_ci_lower"], info["tpr_ci_upper"], alpha=0.2)

    ax1.plot([0, 1], [0, 1], 'k--', label='Random')
    ax1.set_xlabel("False Positive Rate")
    ax1.set_ylabel("True Positive Rate")
    ax1.set_title(auc_title)
    ax1.legend(loc='lower right')
    ax1.grid(True)

    # === 右侧：混淆矩阵 ===
    if cm is not None and ax2 is not None:
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', ax=ax2)
        ax2.set_title(cm_title)
        ax2.set_xlabel("Predicted")
        ax2.set_ylabel("True")

    plt.tight_layout()
    if save_path:
        os.makedirs(save_path, exist_ok=True)
        plt.savefig(os.path.join(save_path, filename))
    plt.close()

def evaluate_model_on_data(model, X, y):
    prob = model.predict_proba(X)[:, 1]
    pred = model.predict(X)
    acc = accuracy_score(y, pred)
    fpr, tpr, _ = roc_curve(y, prob)
    auc_score = auc(fpr, tpr)
    ci_bounds = bootstrap_auc_ci(model, X, y,n_bootstraps=200, ci=0.95, random_state=42)

    return {
        "fpr": fpr,
        "tpr": tpr,
        "auc": auc_score,
        "ci": ci_bounds,
        "acc": acc,
        "prob": prob,
        "pred": pred,
    }
def get_test_metrics(y_true, y_pred,auc_score):
    acc = accuracy_score(y_true, y_pred)
    f1 = f1_score(y_true, y_pred)
    precision = precision_score(y_true, y_pred)
    recall = recall_score(y_true, y_pred)  # Sensitivity

    cm = confusion_matrix(y_true, y_pred)
    tn, fp, fn, tp = cm.ravel()
    specificity = tn / (tn + fp) if (tn + fp) > 0 else 0

    return {
        "Accuracy": acc,
        "F1 Score": f1,
        "Precision": precision,
        "Recall (Sensitivity)": recall,
        "Specificity": specificity,
        "AUC": auc_score,
        "Confusion Matrix": cm
    }


def plot_single_model_train_and_test(model, x_train, y_train, x_test, y_test,
                                     save_path=None, draw_ci=False,
                                     auc_title="ROC Curve", cm_title="Confusion Matrix"):

    result = {'train': evaluate_model_on_data(model, x_train, y_train),
              'test': evaluate_model_on_data(model, x_test, y_test)}
    metrics = get_test_metrics(y_test,result['test']['pred'],result['test']['auc'])
    fw(opj(save_path,'metrics'),format_metrics(metrics))
    cm = metrics['Confusion Matrix']

    plot_auc_cm(
        result=result,
        cm=cm,
        draw_ci=draw_ci,
        auc_title=auc_title,
        cm_title=cm_title,
        save_path=save_path,
        filename="auc_cm.pdf"
    )

def plot_mutil_model_roc(result,
                           save_path=None,filename=None,draw_ci=False,
                           auc_title="ROC Curve"):
    plot_auc_cm(
        result=result,
        draw_ci=draw_ci,
        auc_title=auc_title,
        save_path=save_path,
        filename=filename
    )


def draw_violin_by_set(df, y_col, y_label, title, save_path):
    # 将 0/1 转换为标签
    df = df.copy()
    df['ALN_label'] = df['ALN status'].map({0: 'Non-ALNM', 1: 'ALNM'})

    plt.figure(figsize=(8, 6))
    sns.violinplot(
        x='Dataset',       # Train / Test
        y=y_col,           # 连续变量
        hue='ALN_label',   # 转移状态
        data=df,
        split=True,
        palette='Set2',
        inner='quartile'
    )
    plt.title(title)
    plt.xlabel("")
    plt.ylabel(y_label)
    plt.grid(True, linestyle='--', alpha=0.5)
    plt.legend(title="ALN Status", loc="upper right")
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()


def draw_box_with_centered_points(df, y_col, y_label, title, save_path,p_value=None):
    df = df.copy()
    df['ALN_label'] = df['ALN status'].map({0: 'Non-ALNM', 1: 'ALNM'})

    plt.figure(figsize=(8,6))
    ax = sns.boxplot(
        x='Dataset',
        y=y_col,
        hue='ALN_label',
        data=df,
        palette='Set2',
        fliersize=0,
        showcaps=True,
        boxprops={'alpha':0.6},
        width=0.6
    )

    # 计算箱体对应的x轴坐标位置
    # 箱子默认在整数位置，hue分两类，Seaborn会自动微调0.2左右
    # 一般位置大致为：
    # Dataset类别数为2 (Train=0, Test=1)
    # hue类别数为2 (NoMet=0, Met=1)
    # 每组箱体横坐标约为： dataset_idx +/- offset (offset大约0.2)

    # 定义映射
    dataset_order = ['Train', 'Test']
    hue_order = ['Non-ALNM', 'ALNM']
    offset = 0.2

    # 对每个类别分别绘制散点，横坐标用固定位置，不dodge
    for i, dataset in enumerate(dataset_order):
        for j, aln_label in enumerate(hue_order):
            sub_df = df[(df['Dataset'] == dataset) & (df['ALN_label'] == aln_label)]
            x_center = i - offset if j == 0 else i + offset  # NoMet箱体左偏，Met箱体右偏

            # x轴加随机小抖动避免点完全重叠
            x_jitter = 0
            if j == 0:
                x_jitter = np.random.normal(loc=0.05, scale=0.03, size=len(sub_df))
            else:
                x_jitter = np.random.normal(loc=-0.05, scale=0.03, size=len(sub_df))

            plt.scatter(
                x=np.full(len(sub_df), x_center) + x_jitter,
                y=sub_df[y_col],
                alpha=0.7,
                s=30,
                label=None,
                color=sns.color_palette('Set2')[j]
            )

    # 调整图例：只留箱型图的legend
    handles, labels = ax.get_legend_handles_labels()
    plt.legend(handles, labels, loc='upper right',frameon=False)

    plt.title(title)
    plt.xlabel("")
    plt.ylabel(y_label)
    plt.grid(True, linestyle='--', alpha=0.5)
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()
